library(tidyr)
library(dplyr)
library(stringr)
library(circlize)
library(GenomicRanges)
library(EnrichedHeatmap)
library(ComplexHeatmap)

if (!file.exists("output")) {
  dir.create("output")
}

# load data

source("data_handling.R")
source("HNC_metadata_tidy.R")

# copy number variant information can be found in EGAS00001005450
BAT_DIR <- "/data/mutographs/files/Manuscripts/HNC/Somatic/battenberg/summary/"
list_ids <- unique(str_extract(list.files(BAT_DIR)[str_detect(list.files(BAT_DIR), "battenberg")], "PD[0-9]{5}a"))

cvCOSMIC <- read.csv("Input_signature_attributions/Input_CN_COSMIC_attributions.csv", stringsAsFactors = T) %>% dplyr::rename(donor_id = X)

data <- data %>%
  right_join(cvCOSMIC, by = "donor_id") %>%
  filter(donor_id %in% list_ids) %>%
  arrange(donor_id)

# Calculate categorical and relative frequency
cvM <- tibble::column_to_rownames(cvCOSMIC, "donor_id")
cvM <- t(apply(cvM, 1, function(x) x / sum(x)))
cvM[is.na(cvM)] <- 0

data <- dichotomizeSigs(data, cvCOSMIC)
data <- data %>% left_join(as.data.frame(cvM) %>% tibble::rownames_to_column("donor_id"), by = "donor_id", suffix = c("", "_freq"))

# Generate GRanges object with 22 ranges and 0 metadata column
chr_df <- read.chromInfo()$df
chr_df <- chr_df[chr_df$chr %in% paste0("chr", 1:22), ]
chr_gr <- GRanges(seqnames = chr_df[, 1], ranges = IRanges(chr_df[, 2] + 1, chr_df[, 3]))

# In the final heatmap, each row (if the genomic direction is vertical) or each column (if the genomic direction is horizontal) represents a genomic window
# thus we need to split the genome with equal-width windows.
# EnrichedHeatmap::makeWindows() function is used to split the genome by 1MB window
chr_window <- makeWindows(chr_gr, w = 1e6)

# To visualize genome-scale signals as a heatmap as well as other tracks,
# the average signals in the 1MB windows are calculated by overlapping the genomic windows and the genomic signals.
# This function is adapted from HilbertCurve package following recommendations from ComplexHeatmap recommendations
average_in_window <- function(window, gr, v, method = "weighted", empty_v = NA) {
  if (missing(v)) v <- rep(1, length(gr))
  if (is.null(v)) v <- rep(1, length(gr))
  if (is.atomic(v) && is.vector(v)) v <- cbind(v)

  v <- as.matrix(v)
  if (is.character(v) && ncol(v) > 1) {
    stop("`v` can only be a character vector.")
  }

  if (length(empty_v) == 1) {
    empty_v <- rep(empty_v, ncol(v))
  }

  u <- matrix(rep(empty_v, each = length(window)), nrow = length(window), ncol = ncol(v))

  mtch <- as.matrix(findOverlaps(window, gr))
  intersect <- pintersect(window[mtch[, 1]], gr[mtch[, 2]])
  w <- width(intersect)
  v <- v[mtch[, 2], , drop = FALSE]
  n <- nrow(v)

  ind_list <- split(seq_len(n), mtch[, 1])
  window_index <- as.numeric(names(ind_list))
  window_w <- width(window)

  if (is.character(v)) {
    for (i in seq_along(ind_list)) {
      ind <- ind_list[[i]]
      if (is.function(method)) {
        u[window_index[i], ] <- method(v[ind], w[ind], window_w[i])
      } else {
        tb <- tapply(w[ind], v[ind], sum)
        u[window_index[i], ] <- names(tb[which.max(tb)])
      }
    }
  } else {
    if (method == "w0") {
      gr2 <- reduce(gr, min.gapwidth = 0)
      mtch2 <- as.matrix(findOverlaps(window, gr2))
      intersect2 <- pintersect(window[mtch2[, 1]], gr2[mtch2[, 2]])

      width_intersect <- tapply(width(intersect2), mtch2[, 1], sum)
      ind <- unique(mtch2[, 1])
      width_setdiff <- width(window[ind]) - width_intersect

      w2 <- width(window[ind])

      for (i in seq_along(ind_list)) {
        ind <- ind_list[[i]]
        x <- colSums(v[ind, , drop = FALSE] * w[ind]) / sum(w[ind])
        u[window_index[i], ] <- (x * width_intersect[i] + empty_v * width_setdiff[i]) / w2[i]
      }
    } else if (method == "absolute") {
      for (i in seq_along(ind_list)) {
        u[window_index[i], ] <- colMeans(v[ind_list[[i]], , drop = FALSE])
      }
    } else if (method == "weighted") {
      for (i in seq_along(ind_list)) {
        ind <- ind_list[[i]]
        u[window_index[i], ] <- colSums(v[ind, , drop = FALSE] * w[ind]) / sum(w[ind])
      }
    } else {
      if (is.function(method)) {
        for (i in seq_along(ind_list)) {
          ind <- ind_list[[i]]
          u[window_index[i], ] <- method(v[ind], w[ind], window_w[i])
        }
      } else {
        stop("wrong method.")
      }
    }
  }

  return(u)
}

# Numeric matrix of CNA
num_mat <- NULL
meta <- NULL
for (i in 1:NROW(list_ids)) {
  gr2 <- GRanges_import(list_ids[i], BAT_DIR)
  m <- battenberg_import(list_ids[i], BAT_DIR) %>% mutate(donor_id = list_ids[i])
  num_mat <- cbind(num_mat, average_in_window(chr_window, gr2, gr2$total_cn))
  meta <- rbind(meta, m)
}

colnames(num_mat) <- list_ids

# generate HNC clusters
dist_mat <- dist(t(num_mat), method = "euclidean")
hc_mat <- hclust(dist_mat, method = "ward.D")

CN_clusters <- as.data.frame(cutree(hc_mat, k = c(2, 3, 8))) %>%
  mutate(
    k2 = case_when(`2` == 1 ~ "Cluster D", `2` == 2 ~ "Cluster P"),
    # Clusters D and P are further subdivided into 2 groups each
    HNC_clusters = case_when(
      `2` == 1 & `8` == 1 ~ "Cluster D1",
      `2` == 1 & `8` == 4 ~ "Cluster D2",
      `2` == 2 & `3` == 2 ~ "Cluster P1",
      `2` == 2 & `3` == 3 ~ "Cluster P2"
    ), .keep = "unused"
  ) %>%
  tibble::rownames_to_column("donor_id")

write.csv(CN_clusters, "output/CN_clusters.csv", quote = F, row.names = F)

data <- data %>% left_join(CN_clusters, by = "donor_id")

# Heatmap annotations
CNcol <- rev(scales::viridis_pal()(9))
colors_vir <- rev(scales::viridis_pal()(7))

darkcol <- "#8856a7"
lightcol <- "#ff4538"

top_annotation <- HeatmapAnnotation(
  annotation_name_rot = 0,
  "CN clusters" = data$HNC_clusters,
  "CN" = anno_barplot(data[paste0(colnames(cvM), "_freq")], gp = gpar(fill = CNcol, lwd = 0), bar_width = 1, height = unit(1.2, "cm")),
  show_legend = c("CN clusters" = T, "CN" = F),
  col = list(
    "CN clusters" = c("Cluster D1" = "#225ea8", "Cluster D2" = "#41b6c4", "Cluster P1" = "#a1dab4", "Cluster P2" = "#ffffcc")
  ),
  gap = unit(1, "mm"),
  na_col = "#e6e6e6",
  border = TRUE,
  simple_anno_size = unit(0.4, "cm"),
  annotation_name_side = "left",
  annotation_name_gp = gpar(fontsize = 10)
)

bottom_annotation <- HeatmapAnnotation(
  annotation_name_rot = 0,
  "Tobacco" = data$tobacco,
  "Alcohol" = data$alcohol_ever,
  "HPV" = data$hpv_pos,
  show_legend = c(
    "Tobacco" = T, "Alcohol" = T, "HPV" = T
  ),
  col = list(
    "Tobacco" = c("Current smoker" = "#225ea8", "Ex-smoker" = "#7fcdbb", "Never" = "#edf8b1"),
    "Alcohol" = c("Drinker" = "#225ea8", "Non-drinker" = "#ffffcc"),
    "HPV" = c("Negative" = "#ffffcc", "Positive" = "#225ea8")
  ),
  gap = unit(c(0, 0, 1, 0, 0, 1), "mm"),
  na_col = "#e6e6e6",
  border = TRUE,
  simple_anno_size = unit(0.4, "cm"),
  annotation_name_side = "left",
  annotation_name_gp = gpar(fontsize = 10)
)

col_order <- data %>%
  tibble::column_to_rownames("donor_id") %>%
  arrange(
    desc(tobacco_ever), desc(alcohol_ever),
    desc(CN1_freq), desc(CN2_freq), desc(CN20_freq), desc(CN9_freq), desc(CN18_freq), desc(CN12_freq), desc(CN13_freq), desc(CNV_G_freq), desc(CN5_freq)
  )

# Build heatmap
chr <- as.vector(seqnames(chr_window))
chr <- factor(chr, levels = paste0("chr", 1:22))

ht_opt$TITLE_PADDING <- unit(c(4, 4), "points")
HM <- Heatmap(num_mat,
  name = "CN burden", col = colorRamp2(c(0, 2, 4), c("#4176A1", "white", "#BB3737")),
  row_split = chr,
  cluster_rows = FALSE, show_column_dend = F,
  cluster_columns = F, clustering_distance_columns = "euclidean", clustering_method_columns = "ward.D",
  height = unit(8, "cm"),
  column_title = "CN data (n = 242)",
  column_order = rownames(col_order),
  column_split = data$HNC_clusters,
  show_column_names = F,
  row_title_rot = 0, row_title_gp = gpar(fontsize = 8), border = TRUE,
  row_gap = unit(0, "points"),
  top_annotation = top_annotation,
  bottom_annotation = bottom_annotation
)

myLegends <- packLegend(Legend(title = "CN signatures", labels = names(cvCOSMIC)[-1], legend_gp = gpar(fill = CNcol)),
  direction = "horizontal"
)

pdf(file = paste0("output/Figure_7c_", Sys.Date(), ".pdf"), width = 7, height = 10)
draw(HM,
  heatmap_legend_side = "right",
  annotation_legend_side = "right",
  merge_legend = TRUE,
  annotation_legend_list = myLegends
)
dev.off()
